from keras import Input
from keras.layers import Activation

from .base_layer import BaseLayer
from entity.model_describe import ModelDescribe


class ActivationLayer(BaseLayer):
    def persistence(self, model_list_line: list, train_id: int) -> ModelDescribe:
        m = super().persistence(model_list_line, train_id)

        activation = 'relu' if len(model_list_line) < 2 else model_list_line[1]
        m.var1 = activation
        return m

    def transfer(self, model_list_line: list, inputs: Input, models: list = None):
        activation = 'relu' if len(model_list_line) < 2 else model_list_line[1]
        return Activation(activation)(inputs)
